Skip to content

feat(ascend): add prefix-shared attention Ascend C kernel - #340

Open
zhangj1an wants to merge 2 commits into
RL-Align:testfrom
zhangj1an:feat/ascend-prefix-shared-attention
Open

feat(ascend): add prefix-shared attention Ascend C kernel#340
zhangj1an wants to merge 2 commits into
RL-Align:testfrom
zhangj1an:feat/ascend-prefix-shared-attention

Conversation

@zhangj1an

@zhangj1an zhangj1an commented Aug 25, 2026

Copy link
Copy Markdown
Collaborator

Latest Status [25 Aug 2026]

Will fix CI error then mark as ready for review

Summary

Port of the CUDA prefix-shared fused attention (csrc/cuda/attention/prefix_shared_attention.cu, the GRPO decode workload) to an Ascend NPU version:

  • Forward is implemented as an Ascend C (CANN) kernel, _C_npu.prefix_shared_attention_ascend:

    • Layout q [bs, G, len_q, 128], k/v [bs, len_kv, 128] — the G generated responses share one prompt-prefix KV sequence, stored once per batch instead of once per group (same surface as the CUDA PrefixSharedAttentionOp: bf16 only, D=128, non-causal, no key-padding mask, forward-only).
    • Each (bs, g, 64-row query block) is processed end-to-end by one AI-core block, streaming the keys in a fixed 64-key tile order with fp32 online-softmax accumulation (per-row max / sum-exp rescaling per tile — the same flash-style single pass as the CUDA kernel).
    • No split-K and no cross-block summary merge, so the per-row reduction order depends only on len_kv — outputs are batch-invariant: bitwise identical regardless of batch size, batch position, or how many blocks were launched (items are strided across blocks).
    • UB budget ~146 KB < 192 KB. Correctness-first pure vector-unit implementation (one dot product per key); the CUBE Q·Kᵀ path is left for future work, same as the deterministic attention kernel in PR [WS1][kernels] Deterministic attention Ascend C kernel #320.

  • Integration:

    • gtest spec prefix_shared_attention (op class attention, reusing its tolerance contract) with pytorch / cuda / ascend candidates and a GtestPrefixSharedAttentionOp gold that expands the shared K/V over G and reuses NativeAttentionOp.forward_fp32 (non-causal, default scale).
    • scripts/check_operator.py learns --device npu (defensive torch_npu probe) so the ascend candidates run in the gtest harness on NPU hosts.
    • setup.py .asc glob changed to recursive **/*.asc, allowing kernels to follow the CUDA-style directory structure under csrc/ascend/attention/.

Build notes (same pattern as PR #320)

Each .asc source file can define only one PYBIND11_MODULE, because linking multiple sources with Bisheng causes a duplicate PyInit__C_npu error. Therefore pybind registrations are consolidated in csrc/ascend/ops_npu.asc (logp + prefix-shared attention), while individual kernel files contain only the kernel and host forward functions.

Files

Path Status
csrc/ascend/attention/prefix_shared_attention_ascend.asc Ascend C forward kernel (bf16, fp32 online softmax, fixed 64-key tile order, one block per (bs, g, query block)). New.
csrc/ascend/ops_npu.asc Aggregated _C_npu pybind registration file (logp + prefix-shared attention). New.
csrc/ascend/batch_invariant_logp_ascend.asc Only removes PYBIND11_MODULE, which is moved to the aggregated registration file. Kernel logic is unchanged.
rl_engine/kernels/ops/ascend/attention/prefix_shared_attn.py PrefixSharedAttentionAscendOp wrapper, same surface as the CUDA PrefixSharedAttentionOp (op(q, k, v) -> out) with bf16/D=128 validation. New.
rl_engine/kernels/ops/ascend/attention/__init__.py Package exports. New.
rl_engine/kernels/gtest/operator_specs.py Adds the prefix_shared_attention spec (pytorch/cuda/ascend candidates) and the GtestPrefixSharedAttentionOp gold.
rl_engine/kernels/gtest/operator_inputs.py Adds the prefix-shared input maker and shape name.
scripts/check_operator.py --device npu support (defensive torch_npu probe).
tests/test_prefix_shared_attention_ascend.py Ascend unit tests for correctness / batch invariance / validation. New.
tests/test_ws1_gtest_gpu.py Adds prefix_shared_attention to the registered-ops coverage set.
setup.py Changes the .asc glob to **/*.asc.

Test

Build first (same as PR #320):

source /usr/local/Ascend/ascend-toolkit/set_env.sh
export KERNEL_ALIGN_FORCE_ASCEND=1
pip install -e . --no-build-isolation --no-deps

gtest

python scripts/check_operator.py --op prefix_shared_attention --candidate ascend --device npu --dtype bf16 --batch 2 --seq 64
python scripts/check_operator.py --op prefix_shared_attention --candidate pytorch --device npu --dtype bf16 --batch 2 --seq 64
python scripts/check_operator.py --op batch_invariant_logp --candidate ascend --device npu --dtype bf16

pytest

python -m pytest tests/test_prefix_shared_attention_ascend.py -v
python -m pytest tests/test_batch_invariant_logp.py -v

Test results

Environment: 8× Ascend 910, CANN 8.5.1 (Bisheng), torch 2.7.1 + torch_npu 2.7.1.

Test Result
gtest prefix_shared_attention ascend bf16 (forward) ✅ suite=prefix_shared_attention passed=True pass_rate=1.0000
gtest prefix_shared_attention pytorch bf16 (gold wiring) ✅ suite=prefix_shared_attention passed=True pass_rate=1.0000
gtest batch_invariant_logp ascend bf16 (regression) ✅ suite=batch_invariant_logp passed=True pass_rate=1.0000
pytest tests/test_prefix_shared_attention_ascend.py ✅ 14 passed
pytest tests/test_batch_invariant_logp.py (regression) ✅ 44 passed, 42 skipped (non-NPU backends)
gtest: prefix_shared_attention ascend bf16 (raw) 
INFO 08-25 14:19:47 [RL-Kernel]: Successfully linked to precompiled _C_npu.prefix_shared_attention_ascend kernel.
suite=prefix_shared_attention passed=True pass_rate=1.0000
candidate=ascend-prefix_shared_attention backend=ascend passed=True pass_rate=1.0000
  case=prefix_shared_attention-torch.bfloat16-2x32x64x128 output=0 shape=(2, 32, 64, 128) dtype=torch.bfloat16 max_abs=3.89850140e-03 mean_abs=2.17124965e-04 max_rel=9.33131874e-02 tol=(atol=5.000e-02, rtol=2.000e-02) passed=True
gtest: prefix_shared_attention pytorch bf16 (raw) 
suite=prefix_shared_attention passed=True pass_rate=1.0000
candidate=pytorch-prefix_shared_attention backend=pytorch passed=True pass_rate=1.0000
  case=prefix_shared_attention-torch.bfloat16-2x32x64x128 output=0 shape=(2, 32, 64, 128) dtype=torch.bfloat16 max_abs=1.66380405e-02 mean_abs=6.89244305e-04 max_rel=2.99679004e+03 tol=(atol=5.000e-02, rtol=2.000e-02) passed=True
gtest: batch_invariant_logp ascend bf16 regression (raw) 
INFO 08-25 14:20:07 [RL-Kernel]: Successfully linked to precompiled _C_npu.batch_invariant_logp_ascend kernel.
/home/ma-user/work/z84450661/rl-kernel/rl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.py:105: UserWarning: Cannot create tensor with interal format while allow_internel_format=False, tensor will be created with base format. (Triggered internally at build/CMakeFiles/torch_npu.dir/compiler_depend.ts:335.)
  selected_logp = selected_logp.where(valid_mask, torch.zeros_like(selected_logp))
suite=batch_invariant_logp passed=True pass_rate=1.0000
candidate=ascend-batch_invariant_logp backend=ascend passed=True pass_rate=1.0000
  case=batch_invariant_logp-torch.bfloat16-2x16x257 output=0 shape=(2, 16) dtype=torch.float32 max_abs=4.76837158e-07 mean_abs=1.49011612e-08 max_rel=7.54759455e-08 tol=(atol=6.000e-02, rtol=0.000e+00) passed=True
pytest: tests/test_prefix_shared_attention_ascend.py (raw) 
============================= test session starts ==============================
platform linux -- Python 3.10.20, pytest-9.1.1, pluggy-1.6.0
rootdir: /home/ma-user/work/z84450661/rl-kernel
configfile: pyproject.toml
collecting ... collected 14 items

tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionCorrectness::test_basic PASSED [  7%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionCorrectness::test_exact_tiles PASSED [ 14%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionCorrectness::test_decode_window PASSED [ 21%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionCorrectness::test_long_prefix_multi_tile PASSED [ 28%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionCorrectness::test_shared_kv_across_groups PASSED [ 35%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionBatchInvariance::test_batch_size_1_vs_n PASSED [ 42%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionBatchInvariance::test_different_positions_in_batch PASSED [ 50%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionBatchInvariance::test_block_striding PASSED [ 57%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionBatchInvariance::test_repeated_runs_deterministic PASSED [ 64%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionValidation::test_rejects_fp32 PASSED [ 71%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionValidation::test_rejects_fp16 PASSED [ 78%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionValidation::test_rejects_bad_head_dim PASSED [ 85%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionValidation::test_rejects_4d_kv PASSED [ 92%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionValidation::test_rejects_kv_length_mismatch PASSED [100%]

============================== 14 passed in 5.81s ==============================
pytest: tests/test_batch_invariant_logp.py regression (summary) 
================== 44 passed, 42 skipped, 1 warning in 6.37s ===================

(all TestAscend* classes passed; skips are CUDA/Triton-only tests on this NPU host)

Notes

  • Single-device operator; no CP/SP or other communication. Scope is consistent with the CUDA PrefixSharedAttentionOp.
  • The op is not added to the kernel registry dispatch: the CUDA PrefixSharedAttentionOp is consumed directly (e.g. benchmarks/benchmark_attention.py) and is not registered either; the Ascend op follows the same surface.
  • ruff check passes for all modified Python files.

@zhangj1an
feat(ascend): add prefix-shared attention Ascend C kernel
8e94860 
- csrc/ascend/attention/prefix_shared_attention_ascend.asc: Ascend C
  prefix-shared fused attention forward (bf16, fp32 online softmax over
  fixed 64-key tiles, one block per (bs, g, 64-row query block), no
  split-K -> batch-invariant); D=128, non-causal, same surface as the
  CUDA prefix_shared_attention op
- rl_engine/kernels/ops/ascend/attention/prefix_shared_attn.py:
  PrefixSharedAttentionAscendOp wrapper (op(q, k, v) -> out)
- gtest: prefix_shared_attention spec (pytorch/cuda/ascend candidates)
  with GtestPrefixSharedAttentionOp gold; check_operator.py learns
  --device npu
- pybind consolidated in csrc/ascend/ops_npu.asc (single PYBIND11_MODULE;
  batch_invariant_logp_ascend.asc only drops its module block)
- tests/test_prefix_shared_attention_ascend.py: correctness, batch
  invariance, validation
- setup.py: recursive .asc glob
@coderabbitai

coderabbitai Bot commented Aug 25, 2026 •
edited
Loading

Copy link
Copy Markdown

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 98c2071a-36dc-4306-a26f-2db69a457855

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review

Comment @coderabbitai help to get the list of available commands.

@zhangj1an
Merge remote-tracking branch 'refs/remotes/upstream/test' into feat/a…
c70e632 
…scend-prefix-shared-attention

# Conflicts:
#	rl_engine/kernels/gtest/operator_inputs.py
#	rl_engine/kernels/gtest/operator_specs.py
#	setup.py
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@Flink-ddd Flink-ddd Awaiting requested review from Flink-ddd Flink-ddd is a code owner
@inaniloquentee inaniloquentee Awaiting requested review from inaniloquentee inaniloquentee is a code owner
@bitborne bitborne Awaiting requested review from bitborne bitborne is a code owner
@KJLdefeated KJLdefeated Awaiting requested review from KJLdefeated KJLdefeated is a code owner

At least 2 approving reviews are required to merge this pull request.

Assignees

No one assigned

Labels

Ascend

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@zhangj1an @Flink-ddd